package org.colour.amber.interceptor;

import brave.propagation.CurrentTraceContext;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import java.util.Optional;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.HandlerInterceptor;
import org.springframework.web.servlet.ModelAndView;

/**
 * @author jasonfu
 * @date 2024/9/13 15:39
 */
@Slf4j
@Component
public class TraceInterceptor implements HandlerInterceptor {

    private final CurrentTraceContext traceContext;

    public TraceInterceptor(CurrentTraceContext traceContext) {
        this.traceContext = traceContext;
    }

    /**
     * 在处理请求之前执行拦截操作，主要用于向响应头添加traceId和spanId
     *
     * @param request  HTTP请求对象
     * @param response HTTP响应对象
     * @param handler  被拦截的处理器对象
     * @return boolean 表示是否继续执行下一个拦截器或处理器
     * @throws Exception 如果发生异常，则传递给调用者
     */
    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler)
            throws Exception {
        // 记录进入preHandle方法的日志
        log.info("TraceInterceptor|preHandle");

        // 获取trace上下文对象，并将其traceId添加到响应头中
        Optional.ofNullable(traceContext.get())
                .ifPresent(context -> response.setHeader("traceId", context.traceIdString()));

        // 获取trace上下文对象，并将其spanId添加到响应头中
        Optional.ofNullable(traceContext.get())
                .ifPresent(context -> response.setHeader("spanId", context.spanIdString()));

        // 继续执行下一个拦截器或处理器
        return HandlerInterceptor.super.preHandle(request, response, handler);
    }

    @Override
    public void postHandle(
            HttpServletRequest request, HttpServletResponse response, Object handler, ModelAndView modelAndView)
            throws Exception {
        log.info("TraceInterceptor|postHandle");
        HandlerInterceptor.super.postHandle(request, response, handler, modelAndView);
    }

    @Override
    public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler, Exception ex)
            throws Exception {
        log.info("TraceInterceptor|afterCompletion");
        HandlerInterceptor.super.afterCompletion(request, response, handler, ex);
    }
}
